import collections


class Solution(object):
    def subarraysWithKDistinct(self, nums, k):
        n = len(nums)

        def dfs(val):
            d = collections.defaultdict(int)
            left = 0
            right = 0
            ans = 0
            while right < n:
                d[nums[right]] += 1
                while len(d) > val:
                    d[nums[left]] -= 1
                    if d[nums[left]] == 0:
                        d.pop(nums[left])
                    left += 1
                ans += right - left + 1
                right += 1
            return ans

        return dfs(k) - dfs(k - 1)


data = Solution()
nums = [1, 2, 1, 2, 3]
k = 2
print(data.subarraysWithKDistinct(nums, k))

data = Solution()
nums = [1,2,1,3,4]
k = 3
print(data.subarraysWithKDistinct(nums, k))
